# -*- coding: utf-8 -*-

'''
 :author: Rayco Abad-Martín <rayco.abad@gmail.com>
 :author: Elena Marrero-Méndez <elemarez@gmail.com>
'''

from yacf.Backends.C.Visitors.OmpFilters import *
from Backends.MPI.Mutators.Common import AbstractMPIMutator
from Backends.MPI.Mutators.MPI_OmpFor import MPI_OmpFor
from yacf.Frontend.Shortcuts import getCurrentLanguageAst
from yacf.Frontend.SymbolTable import updateSymbolTableFromNode
from yacf.Tools import SourceStorage
from yacf.Tools.Tree import InsertTool, NodeNotFound, ReplaceTool, RemoveTool

c_ast = getCurrentLanguageAst()

class MPI_OmpParallel(AbstractMPIMutator):
    def __init__(self, file_name = 'loopParallel', *args, **kwargs):
        self.file_name = file_name
        super(MPI_OmpParallel, self).__init__(*args, **kwargs)

    def filter(self, ast):
        """ Filter definition
            Returns the first node matching with the filter
        """
        # Build a visitor , matching the OmpParallel node of the AST
        f = OmpParallelFilter(device = self.device)
        node = f.apply(ast)
        self._func_def = f.get_func_def()
        self._parallel = node
        return node

    def apply_all(self, ast):
        """ Apply mutation to all matches """
        self.ast = ast
        f = OmpParallelFilter(device = self.device)
        num = 0
        old_parallel = self._parallel
        try:
            for elem in f.iterate(ast):
                # Save previous state
                old_file_name = self.file_name
                self.file_name = self.file_name + str(num)
                self._clauses = {}
                # Current scope variables
                self._func_def = f.get_func_def()
                self._target_device_node = f._target_device_node
                # Solve the error when the pragma parallel has not clauses
                if f._target_device_node.clauses == None:
                    f._target_device_node.clauses = {}
                self._clauses = self._get_dict_from_clauses(f._target_device_node.clauses, ast)
                self._parallel = elem
                self.mutatorFunction(ast, elem)
                # Restore previous state
                self.file_name = old_file_name
                self._parallel = old_parallel
                num += 1
        except NodeNotFound as nf:
            print str(nf)
        except StopIteration:
            return self._parallel
        return ast

    def mutatorFunction(self, ast, ompParallel_node):
        # Solve the error when the pragma parallel has not clauses
        if ompParallel_node.clauses == None:
            ompParallel_node.clauses = {}
        clause_dict = self._get_dict_from_clauses(ompParallel_node.clauses, ast)
        shared_params = clause_dict['SHARED']
        modified_shared_vars = clause_dict['COPY_OUT']
        copyin_shared_vars = clause_dict['COPY_IN']
        # If no copy clauses specified, default to shared params
        if len(copyin_shared_vars) == 0 and len(modified_shared_vars) == 0:
            modified_shared_vars = shared_params
            copyin_shared_vars = shared_params
        private_params = clause_dict['PRIVATE'] 
        nowait = clause_dict.has_key('NOWAIT')
        MPI_OmpFor(str(self.file_name), clause_dict, source_storage = self.source_storage).apply_all(ompParallel_node, ast)
        # Statement for MPI
        mpi_stmts = c_ast.Compound(block_items = []);
        InsertTool(subtree = self._parallel.stmt, position = "begin").apply(mpi_stmts, 'block_items')
        ReplaceTool(new_node = mpi_stmts, old_node = self._parallel.parent).apply(self._parallel.parent.parent, 'block_items')
        updateSymbolTableFromNode(ast)

